from scipy.spatial import distance
import numpy as np
import matplotlib.pyplot as plt


class TWVRPModel:
    # 初始化
    def __init__(self, _job_list, _transport_cost_matrix,
                 _company_info, _fix_sequence, _slide_sequence):
        # 任务信息
        self.job_list = _job_list
        # 交通时间信息
        self.transport_cost_matrix = _transport_cost_matrix
        # 工作安排顺序信息
        self.company_info = _company_info
        # 不可移动部分
        self.fix_sequence = _fix_sequence
        # 可移动部分
        self.slide_sequence = _slide_sequence
        # 任务序列实体类
        self.job_sequence_entity_list = []
        # 将公司起点作为第一个任务加入序列实体类
        first_job_entity = {
            # 基本信息
            'job_id': -1,
            'job_name': self.company_info['company_name'],
            'prev_job_id': -1,
            'prev_job_name': self.company_info['company_name'],
            'work_cost': 0,
            'prev_to_now_work_trans_start_time': 0,
            'prev_to_now_work_trans_end_time': 0,
            'now_work_start_time': 0,
            'now_work_end_time': 0,
            # 与位置相关，不需要计算，主要用来画图
            'lon': self.company_info['lon'],
            'lat': self.company_info['lat'],
            'job_location': self.company_info['company_name']
        }
        self.job_sequence_entity_list.append(first_job_entity)

    # 安排可移动任务，固定任务不用安排
    def eachSlideJobAssign(self, _this_job_id):
        # 当前任务的id和名称
        this_job_id = _this_job_id
        # 获取前序任务
        prev_work_id = self.job_sequence_entity_list[-1]['job_id']
        prev_work_end_time = self.job_sequence_entity_list[-1]['now_work_end_time']
        # 将fix_sequence展成一维数组
        fix_sequence_flat = [b for a in self.fix_sequence for b in a]
        # 将起点插入
        fix_sequence_flat.insert(0, -1)
        # 找到job_sequence_entity_list最后一个不可移动节点，默认从起点开始，生成后续序列从而判断是否可以插入
        job_sequence_last_fix_job = -1
        for fix_job in self.job_sequence_entity_list:
            job_id = fix_job['job_id']
            job_class = self.job_list[job_id]['job_class']
            if job_class == 2:
                job_sequence_last_fix_job = job_id
        prev_fix_work_id_index = fix_sequence_flat.index(job_sequence_last_fix_job)
        fix_follow_sequence = fix_sequence_flat[prev_fix_work_id_index + 1:]
        # 判断是否可以插入当前任务
        ifLegal = self.calIfFollowSequenceIflegal(prev_work_id, prev_work_end_time, this_job_id, fix_follow_sequence)
        if ifLegal == 0:
            pass
        else:
            this_job_entity = {
                # 基本信息
                'job_id': this_job_id,
                'job_name': self.job_list[this_job_id]['job_name'],
                'prev_job_id': prev_work_id,
                'prev_job_name': self.job_list[prev_work_id]['job_name'],
                'work_cost': self.job_list[this_job_id]['work_cost'],
                'prev_to_now_work_trans_start_time': prev_work_end_time,
                'prev_to_now_work_trans_end_time': prev_work_end_time + self.transport_cost_matrix[prev_work_id][
                    this_job_id],
                'now_work_start_time': prev_work_end_time + self.transport_cost_matrix[prev_work_id][
                    this_job_id],
                'now_work_end_time': prev_work_end_time + self.transport_cost_matrix[prev_work_id][
                    this_job_id] + self.job_list[this_job_id]['work_cost'],
                # 与位置相关，不需要计算，主要用来画图
                'lon': self.job_list[this_job_id]['lon'],
                'lat': self.job_list[this_job_id]['lat'],
                'job_location': self.job_list[this_job_id]['job_name']
            }
            self.job_sequence_entity_list.append(this_job_entity)
            self.slide_sequence.remove(this_job_id)

    # 计算后续任务是否都合法
    def calIfFollowSequenceIflegal(self, prev_work_id, prev_work_end_time, insert_work_id, follow_sequence):
        ifLegal = 0
        # 插入任务的前序任务到插入任务的距离消耗
        prev_work_insert_work_dist = self.transport_cost_matrix[prev_work_id][insert_work_id]
        insert_work_end_time = prev_work_end_time + \
                               prev_work_insert_work_dist + \
                               self.job_list[insert_work_id]['work_cost']
        # 临时的两个变量，每次都要更新
        sequence_end_time = insert_work_end_time
        sequence_end_work_id = insert_work_id
        for now_work_id in follow_sequence:
            insert_work_now_work_dist = self.transport_cost_matrix[sequence_end_work_id][now_work_id]
            now_work_can_start_time = sequence_end_time + insert_work_now_work_dist
            now_work_start_time = self.job_list[now_work_id]['start_time']
            # 实际开工时间
            now_work_true_start_time = max(now_work_can_start_time, now_work_start_time)
            now_work_work_cost = self.job_list[now_work_id]['work_cost']
            # 实际结束时间
            now_work_true_end_time = now_work_true_start_time + now_work_work_cost
            now_work_end_time = self.job_list[now_work_id]['end_time']
            if now_work_true_end_time > now_work_end_time:
                ifLegal = 0
                break
            else:
                ifLegal = 1
            sequence_end_time = now_work_true_end_time
            sequence_end_work_id = now_work_id
        return ifLegal

    # 全部任务进行安排
    def jobAssign(self):
        # 将fix_sequence展成一维数组
        fix_sequence_flat = [b for a in self.fix_sequence for b in a]
        for fix_sequence_work_id in fix_sequence_flat:
            # 看是否有滑动任务可以插入
            slide_sequence_flat = [a for a in self.slide_sequence]
            for slide_sequence_work_id in slide_sequence_flat:
                self.eachSlideJobAssign(slide_sequence_work_id)

            # 然后加入本个不可移动任务
            prev_work_id = self.job_sequence_entity_list[-1]['job_id']
            prev_work_end_time = self.job_sequence_entity_list[-1]['now_work_end_time']
            this_job_entity = {
                # 基本信息
                'job_id': fix_sequence_work_id,
                'job_name': self.job_list[fix_sequence_work_id]['job_name'],
                'prev_job_id': prev_work_id,
                'prev_job_name': self.job_list[prev_work_id]['job_name'],
                'work_cost': self.job_list[fix_sequence_work_id]['work_cost'],
                'prev_to_now_work_trans_start_time': prev_work_end_time,
                'prev_to_now_work_trans_end_time': prev_work_end_time + self.transport_cost_matrix[prev_work_id][
                    fix_sequence_work_id],
                'now_work_start_time': prev_work_end_time + self.transport_cost_matrix[prev_work_id][
                    fix_sequence_work_id],
                'now_work_end_time': prev_work_end_time + self.transport_cost_matrix[prev_work_id][
                    fix_sequence_work_id] + self.job_list[fix_sequence_work_id]['work_cost'],
                # 与位置相关，不需要计算，主要用来画图
                'lon': self.job_list[fix_sequence_work_id]['lon'],
                'lat': self.job_list[fix_sequence_work_id]['lat'],
                'job_location': self.job_list[fix_sequence_work_id]['job_name']
            }
            self.job_sequence_entity_list.append(this_job_entity)

        # 如果最后可移动任务部分还有，则顺序的接在最后
        slide_sequence_flat = [a for a in self.slide_sequence]
        for slide_sequence_id in slide_sequence_flat:
            # 然后加入本个可移动任务
            prev_work_id = self.job_sequence_entity_list[-1]['job_id']
            prev_work_end_time = self.job_sequence_entity_list[-1]['now_work_end_time']
            this_job_entity = {
                # 基本信息
                'job_id': slide_sequence_id,
                'job_name': self.job_list[slide_sequence_id]['job_name'],
                'prev_job_id': prev_work_id,
                'prev_job_name': self.job_list[prev_work_id]['job_name'],
                'work_cost': self.job_list[slide_sequence_id]['work_cost'],
                'prev_to_now_work_trans_start_time': prev_work_end_time,
                'prev_to_now_work_trans_end_time': prev_work_end_time + self.transport_cost_matrix[prev_work_id][
                    slide_sequence_id],
                'now_work_start_time': prev_work_end_time + self.transport_cost_matrix[prev_work_id][
                    slide_sequence_id],
                'now_work_end_time': prev_work_end_time + self.transport_cost_matrix[prev_work_id][
                    slide_sequence_id] + self.job_list[slide_sequence_id]['work_cost'],
                # 与位置相关，不需要计算，主要用来画图
                'lon': self.job_list[slide_sequence_id]['lon'],
                'lat': self.job_list[slide_sequence_id]['lat'],
                'job_location': self.job_list[slide_sequence_id]['job_name']
            }
            self.job_sequence_entity_list.append(this_job_entity)

    # 获取通线时间(即全部任务完成的时间)
    def getSequenceCostTotalTime(self):
        max_date = self.job_sequence_entity_list[-1]['now_work_end_time']
        return round(max_date,3)

    # 获取通线距离(即全部任务完成的距离)
    def getTotalDistance(self):
        total_distance = 0.0
        # 再画各个任务顺序和任务地点
        for index, value in enumerate(self.job_sequence_entity_list):
            if index == 0:
                pass
            else:
                prev_work_id = self.job_sequence_entity_list[index - 1]['job_id']
                now_work_id = value['job_id']
                mdist = self.transport_cost_matrix[prev_work_id][now_work_id]
                total_distance = total_distance + mdist
        return round(total_distance,3)

    # 画甘特图
    def drawGantt(self):
        MColor = ['red', 'blue', 'yellow', 'orange', 'green', 'moccasin', 'purple',
                  'pink', 'navajowhite', 'Thistle', 'Magenta', 'SlateBlue', 'RoyalBlue',
                  'Aqua', 'floralwhite', 'ghostwhite', 'goldenrod', 'mediumslateblue',
                  'navajowhite', 'navy', 'sandybrown']
        fig = plt.figure(figsize=(15, 6), dpi=70)
        # 横众轴显示字体为‘仿宋’的中文标签
        plt.rcParams['font.sans-serif'] = ['FangSong']
        plt.rcParams['axes.unicode_minus'] = False
        for index, each_job in enumerate(self.job_sequence_entity_list):
            job_id = each_job['job_id']
            job_name = each_job['job_name']
            prev_job_id = each_job['prev_job_id'],
            prev_job_name = each_job['prev_job_name']
            # 计算交通开销
            transport_cost = 0
            transport_cost_text = 'abc'
            # 如果上个节点是公司，则计算从开始转移到中序（公司）节点，再从公司到本任务节点的时间

            transport_cost = each_job['prev_to_now_work_trans_end_time'] \
                             - each_job['prev_to_now_work_trans_start_time']
            transport_cost_text = prev_job_name + '\n' + job_name

            # 画交通时间开销
            plt.barh(job_name, width=transport_cost,
                     height=0.8,
                     color='white',
                     edgecolor='black',
                     left=each_job['prev_to_now_work_trans_start_time'])
            # 画工作开销
            plt.barh(job_name, width=each_job['work_cost'],
                     height=0.8,
                     color=MColor[job_id],
                     edgecolor='black',
                     left=each_job['now_work_start_time'])
            # 交通开销表标明文字
            plt.text(x=each_job['prev_to_now_work_trans_start_time'] + (transport_cost / 2 - 0.25),
                     y=index - 0.2,
                     color='black',
                     weight='bold',
                     s=transport_cost_text,
                     size=8, fontproperties='SimSun')
            # 工作开销标明文字
            plt.text(x=each_job['now_work_start_time'] + (each_job['work_cost'] / 2 - 0.25),
                     y=index - 0.2,
                     color='black',
                     weight='bold',
                     s=job_name,
                     size=12, fontproperties='SimSun')

        plt.ylabel("维修任务", size=20, fontproperties='SimSun')
        plt.xlabel("工作小时", size=20, fontproperties='SimSun')
        plt.tick_params(labelsize=20)
        plt.tick_params(direction='in')
        total_max_time = self.getSequenceCostTotalTime()
        plt.title('通线工时:' + str(total_max_time))
        plt.show()

    # 画位置图
    def drawPlot(self):
        plt.rcParams['font.sans-serif'] = ['FangSong']
        plt.rcParams['axes.unicode_minus'] = False
        # 在标出所有节点的位置
        point_location_x = []
        point_location_y = []
        for job in self.job_sequence_entity_list:
            x = job['lon']
            y = job['lat']
            label = job['job_name']
            point_location_x.append(x)
            point_location_y.append(y)
            plt.annotate(label, (x, y))
        plt.scatter(point_location_x, point_location_y, c='y')

        # 再画各个任务顺序和任务地点
        for index, value in enumerate(self.job_sequence_entity_list):
            if index == 0:
                pass
            else:
                plot_x_set = []
                plot_y_set = []
                # 先添加上一个任务
                prev_work = self.job_sequence_entity_list[index - 1]
                plot_x_set.append(prev_work['lon'])
                plot_y_set.append(prev_work['lat'])
                # 再添加自己
                x = value['lon']
                y = value['lat']
                plot_x_set.append(x)
                plot_y_set.append(y)
                label = str(index + 1) + ":" + prev_work['job_name'] + '-' + value['job_name']
                plt.annotate(label, (((x + prev_work['lon']) / 2 - 0.1), ((y + prev_work['lat']) / 2 - 0.1)))
                plt.plot(plot_x_set, plot_y_set, 'b')

        total_distance = self.getTotalDistance()
        plt.title('总距离' + str(total_distance))
        plt.show()


def create_data():
    # 任务信息,wokr_cost时间开销，单位小时；job_class任务类别，
    # 1为优先任务，2为时间窗任务
    job_list = [
        {'job_id': 0, 'job_name': '1号任务',
         'job_location': 'A地', 'lon': 6, "lat": 6,
         'work_cost': 2.5, 'job_class': 2, 'start_time': 8,
         'end_time': 16},
        {'job_id': 1, 'job_name': '2号任务',
         'job_location': 'B地', 'lon': 8, "lat": 7,
         'work_cost': 2.5, 'job_class': 2, 'start_time': 8,
         'end_time': 16},
        {'job_id': 2, 'job_name': '3号任务',
         'job_location': 'C地', 'lon': 2, "lat": 2,
         'work_cost': 2.5, 'job_class': 2, 'start_time': 24,
         'end_time': 32},
        {'job_id': 3, 'job_name': '4号任务',
         'job_location': 'D地', 'lon': 1, "lat": 4,
         'work_cost': 2.5, 'job_class': 2, 'start_time': 32,
         'end_time': 40},
        {'job_id': 4, 'job_name': '5号任务',
         'job_location': 'E地', 'lon': 9, "lat": 3,
         'work_cost': 2.5, 'job_class': 2, 'start_time': 32,
         'end_time': 40},
        {'job_id': 5, 'job_name': '6号任务',
         'job_location': 'F地', 'lon': 7, "lat": 2,
         'work_cost': 1, 'job_class': 1, 'start_time': None,
         'end_time': None},
        {'job_id': 6, 'job_name': '7号任务',
         'job_location': 'G地', 'lon': 5, "lat": 3,
         'work_cost': 1.5, 'job_class': 1, 'start_time': None,
         'end_time': None},
    ]
    company_info = {
        'company_name': '公司', 'lon': 4, "lat": 4
    }
    # 固定生成8*8的交通开销矩阵，默认第一行第二列为为第一个地点到第二个地点的时间开销，
    # 最后一行为公司到各个地点的开销，这样保证序号都是对的上，每行最后一列表示该点到公司的开销
    transport_cost_matrix = []
    points_list = []
    for index, value in enumerate(job_list):
        points_list.append((value['lon'], value['lat']))
    # 最后还要加上公司
    points_list.append((company_info['lon'], company_info['lat']))
    transport_cost_matrix = distance.cdist(points_list, points_list, 'euclidean')
    # 由于数字太大，全部除以2,且保留一位小数，比如2.2表示需要2.2小时即2小时12分
    transport_cost_matrix = np.around(transport_cost_matrix / 3, decimals=1)
    return job_list, transport_cost_matrix, company_info


if __name__ == '__main__':
    job_list, transport_cost_matrix, company_info = create_data()
    fix_sequence = [[0, 1], [2], [4, 3]]
    slide_sequence = [6, 5]

    twvrpModel = TWVRPModel(job_list, transport_cost_matrix, company_info, fix_sequence, slide_sequence)
    twvrpModel.jobAssign()
    twvrpModel.drawGantt()
    twvrpModel.drawPlot()

    pass
